package mr

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"sort"
	"strconv"
	"time"
)
import "log"
import "net/rpc"
import "hash/fnv"

// Map functions return a slice of KeyValue.
type KeyValue struct {
	Key   string
	Value string
}

// sorting by key
type ByKey []KeyValue

func (a ByKey) Len() int           { return len(a) }
func (a ByKey) Swap(i, j int)      { a[i], a[j] = a[i], a[i] }
func (a ByKey) Less(i, j int) bool { return a[i].Key < a[j].Key }

// use ihash(key) % NReduce to choose the reduce
// task number for each KeyValue emitted by Map.
func ihash(key string) int {
	h := fnv.New32a()
	h.Write([]byte(key))
	return int(h.Sum32() & 0x7fffffff)
}

// main/mrworker.go calls this function.
func Worker(mapf func(string, string) []KeyValue,
	reducef func(string, []string) string) {

	// Your worker implementation here.
	for {
		// 发送请求RPC消息，等待回复消息
		replyMsg := CallForTask()

		switch replyMsg.MsgType {
		case MapTaskAlloc:
			err := HandleMapTask(replyMsg, mapf)

			// 把worker上的执行结果通过RPC发送消息，不需要处理回复消息
			if err == nil {
				_ = CallForReportStatus(MapSuccess, replyMsg.TaskId)
			} else {
				_ = CallForReportStatus(MapFailed, replyMsg.TaskId)
			}
		case ReduceTaskAlloc:
			err := HandleReduceTask(replyMsg, reducef)

			if err == nil {
				_ = CallForReportStatus(ReduceSuccess, replyMsg.TaskId)
			} else {
				_ = CallForReportStatus(ReduceFailed, replyMsg.TaskId)
			}
		case Wait:
			time.Sleep(time.Second * 10) // ten seconds hasn't completed its task
		case Shutdown:
			os.Exit(0)
		default:
			panic("unhandled default case")
		}
		time.Sleep(time.Second)
	}
	// uncomment to send the Example RPC to the coordinator.
	// CallExample()

}

// handle map task
func HandleMapTask(reply *MessageReply, mapf func(string, string) []KeyValue) error {
	// open file
	file, err := os.Open(reply.TaskName)
	if err != nil {
		return err
	}
	defer file.Close()

	// file contents
	content, err := io.ReadAll(file)
	if err != nil {
		return err
	}

	// map function return key-value
	kva := mapf(reply.TaskName, string(content))
	sort.Sort(ByKey(kva))

	// split the file to NReduce files
	tmpFiles := make([]*os.File, reply.NReduce)
	encoders := make([]*json.Encoder, reply.NReduce)

	// kv to split files
	for _, kv := range kva {
		// hash id
		readId := ihash(kv.Key) % reply.NReduce
		if encoders[readId] == nil {
			tmpFile, err := ioutil.TempFile("", fmt.Sprintf("mr-map-tmp-%d", readId))
			if err != nil {
				return err
			}
			defer tmpFile.Close()

			tmpFiles[readId] = tmpFile
			encoders[readId] = json.NewEncoder(tmpFile)
		}

		err := encoders[readId].Encode(&kv)
		if err != nil {
			return err
		}
	}

	// rename the split files
	for i, file := range tmpFiles {
		if file != nil {
			fileName := file.Name()
			file.Close()
			newName := fmt.Sprintf("mr-out-%d-%d", reply.TaskId, i)
			if err := os.Rename(fileName, newName); err != nil {
				return err
			}
		}
	}

	return nil
}

// handle reduce task
func HandleReduceTask(reply *MessageReply, reducef func(string, []string) string) error {
	readid := reply.TaskId
	kva := map[string][]string{}

	// read files from map function
	fileList, err := ReadSpecificFile(readid, "./")
	if err != nil {
		return err
	}

	// put file's key-value to kva
	for _, file := range fileList {
		dec := json.NewDecoder(file)
		for {
			var kv KeyValue
			if err := dec.Decode(&kv); err != nil {
				break
			}
			kva[kv.Key] = append(kva[kv.Key], kv.Value)
		}
		file.Close()
	}

	// sort all keys， reduce函数的key，相当于map函数的value
	var keys []string
	for k := range kva {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	// new a output file
	oname := "mr-out-" + strconv.Itoa(reply.TaskId)
	ofile, err := os.Create(oname)
	if err != nil {
		return nil
	}
	defer ofile.Close()

	// put reduce function output to file
	for _, key := range keys {
		output := reducef(key, kva[key])
		_, err := fmt.Fprintf(ofile, "%v %v\n", key, output)
		if err != nil {
			return err
		}
	}

	DelFileByReduceId(reply.TaskId, "./")

	return nil
}

// the RPC argument and reply are defined in rpc.go
func CallForTask() *MessageReply {
	// argument structure
	args := MessageSend{
		MsgType: AskForTask,
	}

	// reply structure
	reply := MessageReply{}

	// send the RPC request, wait for the reply
	err := call("Coordinator.AskForTask", &args, &reply)

	if err == nil {
		return &reply
	} else {
		return nil
	}
}

// report the task status
func CallForReportStatus(successType MsgType, taskid int) error {
	args := MessageSend{
		MsgType: successType,
		TaskId:  taskid,
	}

	return call("Coordinator.NoticeResult", &args, nil)
}

// example function to show how to make an RPC call to the coordinator.
//
// the RPC argument and reply types are defined in rpc.go.
func CallExample() {

	// declare an argument structure.
	args := ExampleArgs{}

	// fill in the argument(s).
	args.X = 99

	// declare a reply structure.
	reply := ExampleReply{}

	// send the RPC request, wait for the reply.
	// the "Coordinator.Example" tells the
	// receiving server that we'd like to call
	// the Example() method of struct Coordinator.
	ok := call("Coordinator.Example", &args, &reply)
	if ok == nil {
		// reply.Y should be 100.
		fmt.Printf("reply.Y %v\n", reply.Y)
	} else {
		fmt.Printf("call failed!\n")
	}
}

// send an RPC request to the coordinator, wait for the response.
// usually returns no error information.
// returns error information if something goes wrong.
func call(rpcname string, args interface{}, reply interface{}) error {
	// c, err := rpc.DialHTTP("tcp", "127.0.0.1"+":1234")
	sockname := coordinatorSock()
	c, err := rpc.DialHTTP("unix", sockname)
	if err != nil {
		log.Fatal("dialing:", err)
	}
	defer c.Close()

	// invokes the named function, waits for it to complete, and returns its error status
	err = c.Call(rpcname, args, reply)
	//if err == nil {
	//	return true
	//}

	fmt.Println(err)
	return err
}
